# -*- coding: utf-8 -*-
"""
TCN (Temporal Convolutional Network) 预测脚本
使用训练好的TCN模型进行股票价格预测
"""

import pandas as pd
import numpy as np
import sys
import os
import tensorflow as tf
from tensorflow import keras
import joblib
import efinance as ef
import talib
from pypinyin import pinyin, Style
from datetime import datetime

def to_pinyin(text):
    s = ''
    for i in pinyin(text, style=Style.NORMAL):
        s += ''.join(i)
    return s

def configure_gpu_memory():
    """配置计算设备 - 由于RTX 5090兼容性问题，强制使用CPU"""
    # 强制使用CPU，避免RTX 5090兼容性问题
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
    print("🔧 强制使用CPU模式 (RTX 5090兼容性)")
    
    # 优化CPU性能
    import multiprocessing
    cpu_count = multiprocessing.cpu_count()
    tf.config.threading.set_inter_op_parallelism_threads(cpu_count)
    tf.config.threading.set_intra_op_parallelism_threads(cpu_count)
    print(f"⚡ CPU优化: {cpu_count} 线程并行")
    
    return False

def get_latest_stock_data(stock_code, days=200):
    """获取最新的股票数据"""
    try:
        print(f"📊 获取最新 {days} 天的{stock_code}股票数据...")
        
        # 获取历史数据
        df = ef.stock.get_quote_history(stock_code, klt=101, limit=days)
        
        if df.empty:
            raise ValueError("获取数据为空")
        
        # 数据预处理 - 统一列名格式
        column_mapping = {
            '股票代码': 'code', '股票名称': 'name', '日期': 'date',
            '开盘': 'open', '收盘': 'close', '最高': 'high', '最低': 'low', 
            '成交量': 'volume', '成交额': 'amount', '振幅': 'amplitude',
            '涨跌幅': 'change_percent', '涨跌额': 'change_amount', '换手率': 'turnover'
        }
        
        # 同时支持中英文列名
        for old_col, new_col in column_mapping.items():
            if old_col in df.columns:
                df.rename(columns={old_col: new_col}, inplace=True)
                
        # 打印可用列名用于调试
        print(f"Debug: 数据列名: {list(df.columns)}")
        
        # 确保数值类型
        numeric_cols = ['open', 'high', 'low', 'close', 'volume']
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # 设置日期索引
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)
        df.sort_index(inplace=True)
        
        return df
        
    except Exception as e:
        print(f"❌ 获取最新数据失败: {e}")
        return None

def calculate_technical_indicators(df):
    """计算技术指标"""
    df = df.copy()
    
    # 转换为float类型
    close = df['close'].astype(float)
    high = df['high'].astype(float)
    low = df['low'].astype(float)
    volume = df['volume'].astype(float)
    
    # 移动平均线
    df['sma5'] = talib.SMA(close, timeperiod=5)
    df['sma20'] = talib.SMA(close, timeperiod=20)
    
    # 指数移动平均
    df['ema12'] = talib.EMA(close, timeperiod=12)
    df['ema26'] = talib.EMA(close, timeperiod=26)
    
    # MACD
    macd, macdsignal, macdhist = talib.MACD(close, fastperiod=12, slowperiod=26, signalperiod=9)
    df['macd'] = macd
    df['macdsignal'] = macdsignal
    df['macdhist'] = macdhist
    
    # RSI
    df['rsi'] = talib.RSI(close, timeperiod=14)
    
    # 布林带
    upper, middle, lower = talib.BBANDS(close, timeperiod=20)
    df['upperband'] = upper
    df['middleband'] = middle
    df['lowerband'] = lower
    
    # 填充缺失值
    df.ffill(inplace=True)
    df.fillna(0, inplace=True)
    
    return df

def prepare_prediction_data(df, time_steps=60):
    """准备预测数据"""
    feature_columns = [
        'volume', 'sma5', 'sma20', 'ema12', 'ema26', 'macd', 'macdsignal',
        'macdhist', 'rsi', 'upperband', 'middleband', 'lowerband'
    ]
    
    # 获取最新的time_steps条数据作为输入
    features = df[feature_columns].iloc[-time_steps:].values
    
    return features.reshape(1, time_steps, len(feature_columns))

# --- 1. 参数处理 ---
if len(sys.argv) < 3:
    print("错误：参数不足。")
    print("用法: python3 05d_predict_tcn.py <股票代码> <frequency>")
    sys.exit(1)

stock_code = sys.argv[1]
frequency = sys.argv[2]

# --- 2. 获取股票信息 ---
stock_name_pinyin = stock_code
company_name = stock_code
try:
    quote_df = ef.stock.get_quote_history(stock_code, klt=101, end='20300101', limit=2)
    if not quote_df.empty and '股票名称' in quote_df.columns:
        company_name = quote_df['股票名称'].iloc[0]
        stock_name_pinyin = to_pinyin(company_name)
        print(f"📊 股票信息: {company_name} ({stock_code})")
except Exception as e:
    print(f"警告：未能获取股票名称。错误: {e}")

# --- 3. 配置计算环境 ---
use_gpu = configure_gpu_memory()

# --- 4. 定义文件路径 ---
model_path = f'output/models/tcn_model_{stock_code}_{frequency}.h5'
scaler_X_path = f'output/scalers/tcn_X_scaler_{stock_code}_{frequency}.pkl'
scaler_y_path = f'output/scalers/tcn_y_scaler_{stock_code}_{frequency}.pkl'

print(f"""
🔮 TCN股票预测系统
================================
📊 股票代码: {stock_code}
🏢 公司名称: {company_name}
⏰ 预测周期: {frequency}
📅 预测时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
💻 计算模式: {'GPU' if use_gpu else 'CPU'}
""")

# --- 5. 加载模型和归一化器 ---
print("🤖 加载TCN模型和归一化器...")

# 检查文件存在性
price_info_path = f'output/scalers/price_info_{stock_code}_{frequency}.pkl'
required_files = [model_path, scaler_X_path, scaler_y_path, price_info_path]
for file_path in required_files:
    if not os.path.exists(file_path):
        print(f"❌ 未找到必需文件: {file_path}")
        print(f"   请先运行: python3 04d_train_tcn.py {stock_code} {frequency}")
        sys.exit(1)

try:
    # 导入TCN类以便模型加载
    from tcn import TCN
    
    # 加载模型时包含TCN类
    model = keras.models.load_model(model_path, custom_objects={'TCN': TCN})
    print(f"✅ TCN模型加载成功: {model_path}")
    
    # 加载归一化器
    scaler_X = joblib.load(scaler_X_path)
    scaler_y = joblib.load(scaler_y_path)
    
    # 🚀 加载价格信息（用于变化率转换）
    price_info = joblib.load(price_info_path)
    print("✅ 归一化器和价格信息加载成功")
    
except Exception as e:
    print(f"❌ 加载模型失败: {e}")
    sys.exit(1)

# --- 6. 获取最新数据并计算指标 ---
print("📊 获取最新市场数据...")

latest_df = get_latest_stock_data(stock_code, days=200)
if latest_df is None:
    print("❌ 无法获取最新数据，预测失败")
    sys.exit(1)

# 计算技术指标
latest_df = calculate_technical_indicators(latest_df)
print(f"✅ 技术指标计算完成，数据日期范围: {latest_df.index[0].date()} 至 {latest_df.index[-1].date()}")

# 获取最新价格信息
current_data = latest_df.iloc[-1]
current_open = current_data['open']
current_close = current_data['close']
current_high = current_data['high']
current_low = current_data['low']
current_volume = current_data['volume']

# --- 7. 准备预测输入数据 ---
print("🔧 准备预测输入数据...")

TIME_STEPS = 60
prediction_input = prepare_prediction_data(latest_df, TIME_STEPS)

print(f"📊 输入数据形状: {prediction_input.shape}")
print(f"📊 数据日期范围: {latest_df.index[-TIME_STEPS].date()} 至 {latest_df.index[-1].date()}")

# --- 8. 数据归一化 ---
print("🔧 应用归一化...")

# 对输入数据进行归一化
# 将3D数据重塑为2D来应用归一化器
batch_size, time_steps, features = prediction_input.shape
input_reshaped = prediction_input.reshape(-1, features)
input_scaled_reshaped = scaler_X.transform(input_reshaped)
input_scaled = input_scaled_reshaped.reshape(batch_size, time_steps, features)

# --- 9. 模型预测 ---
print("🔮 执行TCN预测...")
print("🚀 使用价格变化率模型进行预测...")

try:
    prediction_scaled = model.predict(input_scaled, verbose=0)
    predicted_change_rates = scaler_y.inverse_transform(prediction_scaled)[0]
    
    # 🚀 关键步骤：从变化率转换为绝对价格
    # 预测价格 = 当前价格 * (1 + 预测变化率)
    predicted_open = current_close * (1 + predicted_change_rates[0])   # 基于前日收盘价预测开盘价
    predicted_close = current_close * (1 + predicted_change_rates[1])  # 基于前日收盘价预测收盘价
    
    print("✅ 预测完成")
    print(f"📊 预测变化率: 开盘 {predicted_change_rates[0]:+.4f} ({predicted_change_rates[0]*100:+.2f}%), 收盘 {predicted_change_rates[1]:+.4f} ({predicted_change_rates[1]*100:+.2f}%)")
    
except Exception as e:
    print(f"❌ 预测失败: {e}")
    sys.exit(1)

# --- 10. 结果分析 ---
open_change = predicted_open - current_close  # 相对于当前收盘价的变化
close_change = predicted_close - current_close
open_change_pct = (open_change / current_close) * 100
close_change_pct = (close_change / current_close) * 100

# 趋势判断
if close_change_pct > 1:
    trend_emoji = "📈"
    trend_text = "强势看涨"
    trend_color = "🟢"
elif close_change_pct > 0:
    trend_emoji = "📈" 
    trend_text = "小幅看涨"
    trend_color = "🟢"
elif close_change_pct > -1:
    trend_emoji = "📉"
    trend_text = "小幅看跌" 
    trend_color = "🟡"
else:
    trend_emoji = "📉"
    trend_text = "明显看跌"
    trend_color = "🔴"

# --- 11. 输出预测结果 ---
print(f"""
{'='*60}
🔮 TCN模型预测结果
{'='*60}
📊 股票代码: {stock_code} ({company_name})
⏰ 预测周期: {frequency}
📅 数据基准: {latest_df.index[-1].date()}
------------------------------------------------------------
💰 当前价格信息:
📈 当日开盘价: {current_open:.2f}元
📊 当日收盘价: {current_close:.2f}元  
📈 当日最高价: {current_high:.2f}元
📉 当日最低价: {current_low:.2f}元
💾 当日成交量: {current_volume:,.0f}手
------------------------------------------------------------
🔮 TCN预测结果:
🌅 预测开盘价: {predicted_open:.2f}元 ({open_change:+.2f}元, {open_change_pct:+.2f}%)
🌇 预测收盘价: {predicted_close:.2f}元 ({close_change:+.2f}元, {close_change_pct:+.2f}%)
📈 预测趋势: {trend_color} {trend_emoji} {trend_text}
{'='*60}

🤖 TCN模型特色:
• Google验证的高性能时间序列预测方法  
• 训练速度比LSTM快10倍+
• 并行计算支持，梯度稳定
• 长期依赖建模能力强

⚠️ 风险提示:
• 本预测仅供参考，不构成投资建议
• 股市有风险，投资需谨慎  
• 请结合其他分析方法综合判断
• 建议设置合理的止盈止损策略
""")

# --- 12. 保存预测记录 (可选) ---
try:
    prediction_record = {
        'timestamp': datetime.now(),
        'stock_code': stock_code,
        'company_name': company_name,
        'current_close': current_close,
        'predicted_open': predicted_open,
        'predicted_close': predicted_close,
        'open_change_pct': open_change_pct,
        'close_change_pct': close_change_pct,
        'trend': trend_text,
        'model': 'TCN'
    }
    
    # 可以将预测记录保存到CSV文件
    # record_file = f'output/predictions/tcn_predictions_{stock_code}.csv'
    # 这里暂时跳过，避免文件过多
    
except Exception as e:
    print(f"⚠️ 保存预测记录时出错: {e}")

print("🚀 TCN预测流程完成！")